import pygame
import math
import sys

# 初始化Pygame
pygame.init()

# 屏幕尺寸
WIDTH, HEIGHT = 1200, 800
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("3D太阳系模拟")

# 设置中文字体
font = pygame.font.SysFont('simhei', 16)  # 使用黑体支持中文
font_big = pygame.font.SysFont('simhei', 24)  # 大号字体

# 颜色定义
BLACK = (0, 0, 0)
YELLOW = (255, 255, 0)  # 太阳
GRAY = (169, 169, 169)  # 水星
LIGHT_BROWN = (210, 180, 140)  # 金星
BLUE = (0, 0, 255)  # 地球
RED = (255, 0, 0)  # 火星
ORANGE = (255, 165, 0)  # 木星
GOLD = (255, 215, 0)  # 土星
CYAN = (0, 255, 255)  # 天王星
DARK_BLUE = (0, 0, 139)  # 海王星
WHITE = (255, 255, 255)


class Planet:
    def __init__(self, name, radius, color, distance, speed, angle=0, tilt=0):
        self.name = name
        self.radius = radius
        self.color = color
        self.distance = distance
        self.speed = speed
        self.angle = angle
        self.tilt = tilt  # 轨道倾角
        self.orbit_points = []
        self.max_orbit_points = 500

    def update(self, time_scale):
        # 更新行星位置
        self.angle += self.speed * time_scale
        if self.angle > 2 * math.pi:
            self.angle -= 2 * math.pi
        elif self.angle < 0:
            self.angle += 2 * math.pi

        # 记录轨道点（3D投影到2D）
        x = WIDTH // 2 + self.distance * math.cos(self.angle)
        y = HEIGHT // 2 + self.distance * math.sin(self.angle) * math.cos(self.tilt)
        z = self.distance * math.sin(self.angle) * math.sin(self.tilt)

        # 3D投影到2D屏幕
        perspective = 1000 / (1000 + z)
        screen_x = WIDTH // 2 + (x - WIDTH // 2) * perspective
        screen_y = HEIGHT // 2 + (y - HEIGHT // 2) * perspective

        self.orbit_points.append((screen_x, screen_y))
        if len(self.orbit_points) > self.max_orbit_points:
            self.orbit_points.pop(0)

    def draw(self, screen):
        # 绘制轨道
        if len(self.orbit_points) > 1:
            pygame.draw.lines(screen, (*self.color, 100), False, self.orbit_points, 1)

        # 计算行星在3D空间中的位置
        x = WIDTH // 2 + self.distance * math.cos(self.angle)
        y = HEIGHT // 2 + self.distance * math.sin(self.angle) * math.cos(self.tilt)
        z = self.distance * math.sin(self.angle) * math.sin(self.tilt)

        # 3D投影到2D屏幕
        perspective = 1000 / (1000 + z)
        screen_x = WIDTH // 2 + (x - WIDTH // 2) * perspective
        screen_y = HEIGHT // 2 + (y - HEIGHT // 2) * perspective
        display_radius = max(1, int(self.radius * perspective))

        # 绘制行星
        pygame.draw.circle(screen, self.color, (int(screen_x), int(screen_y)), display_radius)

        # 绘制行星名称
        name_text = font.render(self.name, True, WHITE)
        screen.blit(name_text, (screen_x - name_text.get_width() // 2, screen_y - display_radius - 20))


class Sun:
    def __init__(self):
        self.radius = 30
        self.color = YELLOW
        self.flare_angle = 0

    def update(self, time_scale):
        self.flare_angle += 0.01 * time_scale
        if self.flare_angle > 2 * math.pi:
            self.flare_angle -= 2 * math.pi

    def draw(self, screen):
        # 绘制太阳
        pygame.draw.circle(screen, self.color, (WIDTH // 2, HEIGHT // 2), self.radius)

        # 绘制太阳耀斑效果
        for i in range(5):
            angle = self.flare_angle + i * (2 * math.pi / 5)
            flare_length = 10 + 5 * math.sin(pygame.time.get_ticks() * 0.005)
            end_x = WIDTH // 2 + (self.radius + flare_length) * math.cos(angle)
            end_y = HEIGHT // 2 + (self.radius + flare_length) * math.sin(angle)
            pygame.draw.line(screen, YELLOW,
                             (WIDTH // 2, HEIGHT // 2),
                             (end_x, end_y), 2)


def create_solar_system():
    # 创建太阳
    sun = Sun()

    # 创建行星: (名称, 半径, 颜色, 距离, 速度, 角度, 轨道倾角)
    planets = [
        Planet("水星", 5, GRAY, 70, 0.02, 0, 0.1),
        Planet("金星", 8, LIGHT_BROWN, 100, 0.015, 0, 0.05),
        Planet("地球", 9, BLUE, 140, 0.01, 0, 0),
        Planet("火星", 7, RED, 180, 0.008, 0, 0.08),
        Planet("木星", 20, ORANGE, 230, 0.005, 0, 0.12),
        Planet("土星", 18, GOLD, 290, 0.004, 0, 0.15),
        Planet("天王星", 12, CYAN, 340, 0.003, 0, 0.2),
        Planet("海王星", 12, DARK_BLUE, 380, 0.002, 0, 0.25)
    ]

    return sun, planets


def main():
    # 创建太阳系
    sun, planets = create_solar_system()

    # 时间缩放因子（控制模拟速度）
    time_scale = 1.0

    # 主循环
    clock = pygame.time.Clock()
    running = True

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_UP:
                    time_scale = min(time_scale + 0.2, 5.0)
                elif event.key == pygame.K_DOWN:
                    time_scale = max(time_scale - 0.2, 0.2)
                elif event.key == pygame.K_r:
                    # 重置太阳系
                    sun, planets = create_solar_system()
                elif event.key == pygame.K_ESCAPE:
                    running = False

        # 填充背景
        screen.fill(BLACK)

        # 绘制星空背景
        for i in range(100):
            x = (i * 17) % WIDTH
            y = (i * 13) % HEIGHT
            size = (i % 3) + 1
            brightness = 100 + (i % 155)
            pygame.draw.circle(screen, (brightness, brightness, brightness), (x, y), size)

        # 更新和绘制太阳
        sun.update(time_scale)
        sun.draw(screen)

        # 更新和绘制行星（按距离排序，实现遮挡效果）
        sorted_planets = sorted(planets, key=lambda p:
        p.distance * math.sin(p.angle) * math.sin(p.tilt))
        for planet in sorted_planets:
            planet.update(time_scale)
            planet.draw(screen)

        # 绘制UI信息
        time_text = font.render(f"速度: {time_scale:.1f}x (↑/↓ 调整)", True, WHITE)
        restart_text = font.render("R: 重置系统  ESC: 退出", True, WHITE)
        title_text = font_big.render("3D太阳系模拟", True, WHITE)
        screen.blit(title_text, (WIDTH // 2 - title_text.get_width() // 2, 10))
        screen.blit(time_text, (10, HEIGHT - 60))
        screen.blit(restart_text, (10, HEIGHT - 35))

        # 更新显示
        pygame.display.flip()
        clock.tick(60)

    pygame.quit()
    sys.exit()


if __name__ == "__main__":
    main()
